from typing import List

from langchain_core.embeddings import Embeddings
from sentence_transformers import SentenceTransformer


class M3eEmbeddings(Embeddings):
    def __init__(self, model: SentenceTransformer):
        super().__init__()
        self.model = model

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """
        使用 self.model 对文档列表进行嵌入，并转换为预期的列表格式。
        """
        embeddings = []
        for text in texts:
            raw_embedding = self.model.encode(text)
            converted_embedding = raw_embedding.tolist()
            embeddings.append(converted_embedding)
        return embeddings

    def embed_query(self, text: str) -> List[float]:
        """
        使用 self.model 对查询文本进行嵌入，并转换为预期的列表格式。
        """
        raw_embedding = self.model.encode(text)
        return raw_embedding.tolist()
